{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G3MMAcssHTML"
      },
      "source": [
        "<link rel=\"stylesheet\" href=\"/site-assets/css/gemma.css\">\n",
        "<link rel=\"stylesheet\" href=\"https://fonts.googleapis.com/css2?family=Google+Symbols:opsz,wght,FILL,GRAD@20..48,100..700,0..1,-50..200\" />"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N_yUpPhqrRrK"
      },
      "source": [
        "# Inference with CodeGemma using JAX and Flax"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-yDXE-RX835U"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://ai.google.dev/gemma/docs/codegemma/codegemma_flax_inference\"><img src=\"https://ai.google.dev/static/site-assets/images/docs/notebook-site-button.png\" height=\"32\" width=\"32\" />View on ai.google.dev</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/gemma/docs/codegemma/codegemma_flax_inference.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/google/generative-ai-docs/blob/main/site/en/gemma/docs/codegemma/codegemma_flax_inference.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MUnQEMHBt3nc"
      },
      "source": [
        "We present CodeGemma, a collection of open code models based on Google DeepMind’s Gemma models (Gemma Team et al., 2024).\n",
        "CodeGemma is a family of lightweight, state-of-the art open models built from the same research and technology used to create the Gemini models.\n",
        "\n",
        "Continuing from Gemma pretrained models, CodeGemma models are further trained on more than 500 to 1000 billion tokens of primarily code, using\n",
        "the same architectures as the Gemma model family. As a result, CodeGemma models achieve state of-the-art code performance in both completion\n",
        "and generation tasks, while maintaining strong\n",
        "understanding and reasoning skills at scale.\n",
        "\n",
        "CodeGemma has 3 variants:\n",
        "\n",
        "* A 7B code pretrained model\n",
        "* A 7B instruction-tuned code model\n",
        "* A 2B model, trained specifically for code infilling and open-ended generation.\n",
        "\n",
        "This guide walks you through using the CodeGemma model with Flax for a code completion task.\n",
        "\n",
        "**Note:** This notebook runs on TPU v2 in Google Colab because T4 GPU has insufficient memory."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dbRLI7Q4-8Ve"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n8Ku4iK6PnC0"
      },
      "source": [
        "### 1. Set up Kaggle access for CodeGemma\n",
        "\n",
        "To complete this tutorial, you first need to follow the setup instructions at [Gemma setup](https://ai.google.dev/gemma/docs/setup), which show you how to do the following:\n",
        "\n",
        "* Get access to CodeGemma on [kaggle.com](https://www.kaggle.com/models/google/codegemma/).\n",
        "* Select a Colab runtime with sufficient resources (**T4 GPU has insufficient memory, use TPU v2 instead**) to run the CodeGemma model.\n",
        "* Generate and configure a Kaggle username and API key.\n",
        "\n",
        "After you've completed the Gemma setup, move on to the next section, where you'll set environment variables for your Colab environment.\n",
        "\n",
        "### 2. Set environment variables\n",
        "\n",
        "Set environment variables for `KAGGLE_USERNAME` and `KAGGLE_KEY`. When prompted with the \"Grant access?\" messages, agree to provide secret access."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AVH6Y4k2964n"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from google.colab import userdata # `userdata` is a Colab API.\n",
        "\n",
        "os.environ[\"KAGGLE_USERNAME\"] = userdata.get('KAGGLE_USERNAME')\n",
        "os.environ[\"KAGGLE_KEY\"] = userdata.get('KAGGLE_KEY')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m1UE1CEnE9ql"
      },
      "source": [
        "### 3. Install the `gemma` library\n",
        "\n",
        "Free Colab hardware acceleration is currently *insufficient* to run this notebook. If you are using [Colab Pay As You Go or Colab Pro](https://colab.research.google.com/signup), click on **Edit** > **Notebook settings** > Select **A100 GPU** > **Save** to enable hardware acceleration.\n",
        "\n",
        "Next, you need to install the Google DeepMind `gemma` library from [`github.com/google-deepmind/gemma`](https://github.com/google-deepmind/gemma). If you get an error about \"pip's dependency resolver\", you can usually ignore it.\n",
        "\n",
        "**Note:** By installing `gemma`, you will also install [`flax`](https://flax.readthedocs.io), core [`jax`](https://jax.readthedocs.io), [`orbax`](https://orbax.readthedocs.io/), and [`sentencepiece`](https://github.com/google/sentencepiece)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XpSw-_4EEcoY"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m133.7/133.7 kB\u001b[0m \u001b[31m1.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for gemma (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ],
      "source": [
        "!pip install -q git+https://github.com/google-deepmind/gemma.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-mRkkT-iPYoq"
      },
      "source": [
        "### 4. Import libraries\n",
        "\n",
        "This notebook uses [Gemma](https://github.com/google-deepmind/gemma) (which uses [Flax](https://flax.readthedocs.io) to build its neural network layers), and [SentencePiece](https://github.com/google/sentencepiece) (for tokenization)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ChMf1H4mPVx_"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from gemma import params as params_lib\n",
        "from gemma import sampler as sampler_lib\n",
        "from gemma import transformer as transformer_lib\n",
        "import sentencepiece as spm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oNgKIkxMOsit"
      },
      "source": [
        "## Load the CodeGemma model\n",
        "\n",
        "Load the CodeGemma model with [`kagglehub.model_download`](https://github.com/Kaggle/kagglehub/blob/bddefc718182282882b72f814d407d89e5d178c4/src/kagglehub/models.py#L12), which takes three arguments:\n",
        "\n",
        "- `handle`: The model handle from Kaggle\n",
        "- `path`: (Optional string) The local path\n",
        "- `force_download`: (Optional boolean) Forces to re-download the model\n",
        "\n",
        "**Note:** Be mindful that the `2b-pt` model is around 3.66Gb in size."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X-i10429N-g2"
      },
      "outputs": [],
      "source": [
        "GEMMA_VARIANT = '2b-pt' # @param ['2b-pt', '7b-it', '7b-pt', '1.1-2b-pt', '1.1-7b-it'] {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j_QdPAGyO5zl"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Warning: Looks like you're using an outdated `kagglehub` version, please consider updating (latest version: 0.2.7)\n",
            "Downloading from https://www.kaggle.com/api/v1/models/google/codegemma/flax/2b-pt/3/download...\n",
            "100%|██████████| 3.67G/3.67G [00:22<00:00, 173MB/s]\n",
            "Extracting model files...\n"
          ]
        }
      ],
      "source": [
        "import kagglehub\n",
        "\n",
        "GEMMA_PATH = kagglehub.model_download(f'google/codegemma/flax/{GEMMA_VARIANT}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cjnXlLkWcHIy"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "GEMMA_PATH: /root/.cache/kagglehub/models/google/codegemma/flax/2b-pt/3\n"
          ]
        }
      ],
      "source": [
        "print('GEMMA_PATH:', GEMMA_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E1HzOpDcM04q"
      },
      "source": [
        "**Note:** The path from the output above is where the model weights and tokenizer are saved locally, you will need them for later."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ytvcJ8FPEMm"
      },
      "source": [
        "Check the location of the model weights and the tokenizer, then set the path variables. The tokenizer directory will be in the main directory where you downloaded the model, while the model weights will be in a sub-directory. For example:\n",
        "\n",
        "- The `spm.model` tokenizer file will be in `/LOCAL/PATH/TO/codegemma/flax/2b-pt/3`\n",
        "- The model checkpoint will be in `/LOCAL/PATH/TO/codegemma/flax/2b-pt/3/2b-pt`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JAwXvpzbuiB5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "CKPT_PATH: /root/.cache/kagglehub/models/google/codegemma/flax/2b-pt/3/2b-pt\n",
            "TOKENIZER_PATH: /root/.cache/kagglehub/models/google/codegemma/flax/2b-pt/3/spm.model\n"
          ]
        }
      ],
      "source": [
        "CKPT_PATH = os.path.join(GEMMA_PATH, GEMMA_VARIANT[-5:])\n",
        "TOKENIZER_PATH = os.path.join(GEMMA_PATH, 'spm.model')\n",
        "print('CKPT_PATH:', CKPT_PATH)\n",
        "print('TOKENIZER_PATH:', TOKENIZER_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jc0ZzYIW0TSN"
      },
      "source": [
        "## Perform sampling/inference"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aEe3p8geqekV"
      },
      "source": [
        "Load and format the CodeGemma model checkpoint with the [`gemma.params.load_and_format_params`](https://github.com/google-deepmind/gemma/blob/c6bd156c246530e1620a7c62de98542a377e3934/gemma/params.py#L27) method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "by6eWKtqzxRf"
      },
      "outputs": [],
      "source": [
        "params = params_lib.load_and_format_params(CKPT_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Xpnb2igrGjk"
      },
      "source": [
        "Load the CodeGemma tokenizer, constructed using [`sentencepiece.SentencePieceProcessor`](https://github.com/google/sentencepiece/blob/4d6a1f41069c4636c51a5590f7578a0dbed83450/python/src/sentencepiece/__init__.py#L423):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TpyG5YW1EcoY"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "execution_count": 10,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "vocab = spm.SentencePieceProcessor()\n",
        "vocab.Load(TOKENIZER_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BtJhJkkZzsy1"
      },
      "source": [
        "To automatically load the correct configuration from the CodeGemma model checkpoint, use [`gemma.transformer.TransformerConfig`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L65). The `cache_size` argument is the number of time steps in the CodeGemma `Transformer` cache. Afterwards, instantiate the CodeGemma model as `model_2b` with [`gemma.transformer.Transformer`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L136) (which inherits from [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html)).\n",
        "\n",
        "**Note:** The vocabulary size is smaller than the number of input embeddings because of unused tokens in the current CodeGemma release."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_jjlFAkazzit"
      },
      "outputs": [],
      "source": [
        "transformer_config = transformer_lib.TransformerConfig.from_params(\n",
        "    params,\n",
        "    cache_size=1024\n",
        ")\n",
        "\n",
        "transformer = transformer_lib.Transformer(config=transformer_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EtfVo3pDDAZV"
      },
      "source": [
        "Create a `sampler` with [`gemma.sampler.Sampler`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/sampler.py#L88). It uses the CodeGemma model checkpoint and the tokenizer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dQ1oCF10Ecod"
      },
      "outputs": [],
      "source": [
        "sampler = sampler_lib.Sampler(\n",
        "    transformer=transformer,\n",
        "    vocab=vocab,\n",
        "    params=params['transformer']\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gOi4ua6axnGD"
      },
      "source": [
        "Create some variables to represent the fill-in-the-middle (fim) tokens and create some helper functions to format the prompt and generated output.\n",
        "\n",
        "For example, let's look at the following code:\n",
        "```\n",
        "def function(string):\n",
        "assert function('asdf') == 'fdsa'\n",
        "```\n",
        "We would like to fill in the `function` so that the assertion holds `True`. In this case, the prefix would be:\n",
        "```\n",
        "\"def function(string):\\n\"\n",
        "```\n",
        "And the suffix would be:\n",
        "```\n",
        "\"assert function('asdf') == 'fdsa'\"\n",
        "```\n",
        "We then format this into a prompt as PREFIX-SUFFIX-MIDDLE (the middle section that needs to be filled is always at the end of the prompt):\n",
        "```\n",
        "\"<|fim_prefix|>def function(string):\\n<|fim_suffix|>assert function('asdf') == 'fdsa'<|fim_middle|>\"\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xgkQgFH1xyP2"
      },
      "outputs": [],
      "source": [
        "# In the context of a code editor,\n",
        "# the cursor is the location where the text will be inserted\n",
        "BEFORE_CURSOR = \"<|fim_prefix|>\"\n",
        "AFTER_CURSOR = \"<|fim_suffix|>\"\n",
        "AT_CURSOR = \"<|fim_middle|>\"\n",
        "FILE_SEPARATOR = \"<|file_separator|>\"\n",
        "\n",
        "def format_completion_prompt(before, after):\n",
        "  print(f\"\\nORIGINAL PROMPT:\\n{before}{after}\")\n",
        "  prompt = f\"{BEFORE_CURSOR}{before}{AFTER_CURSOR}{after}{AT_CURSOR}\"\n",
        "  print(f\"\\nFORMATTED PROMPT:\\n{repr(prompt)}\")\n",
        "  return prompt\n",
        "def format_generated_output(before, after, output):\n",
        "  print(f\"\\nGENERATED OUTPUT:\\n{repr(output)}\")\n",
        "  formatted_output = f\"{before}{output.replace(FILE_SEPARATOR, '')}{after}\"\n",
        "  print(f\"\\nFILL-IN COMPLETION:\\n{formatted_output}\")\n",
        "  return formatted_output"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-61KZz7EHiIS"
      },
      "source": [
        "Create a prompt and perform inference. Specify the prefix `before` text and the suffix `after` text and generate the formatted prompt using the helper function `format_completion prompt`.\n",
        "\n",
        "You can tweak `total_generation_steps` (the number of steps performed when generating a response — this example uses `100` to preserve host memory).\n",
        "\n",
        "**Note:** If you run out of memory, click on **Runtime** > **Disconnect and delete runtime**, and then **Runtime** > **Run all**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S5F3fk22Ecod"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "ORIGINAL PROMPT:\n",
            "def function(string):\n",
            "assert function('asdf') == 'fdsa'\n",
            "\n",
            "FORMATTED PROMPT:\n",
            "\"<|fim_prefix|>def function(string):\\n<|fim_suffix|>assert function('asdf') == 'fdsa'<|fim_middle|>\"\n",
            "\n",
            "GENERATED OUTPUT:\n",
            "'    return string[::-1]\\n\\n<|file_separator|>'\n",
            "\n",
            "FILL-IN COMPLETION:\n",
            "def function(string):\n",
            "    return string[::-1]\n",
            "\n",
            "assert function('asdf') == 'fdsa'\n"
          ]
        }
      ],
      "source": [
        "before = \"def function(string):\\n\"\n",
        "after = \"assert function('asdf') == 'fdsa'\"\n",
        "prompt = format_completion_prompt(before, after)\n",
        "\n",
        "output = sampler(\n",
        "    [prompt],\n",
        "    total_generation_steps=100,\n",
        "    ).text\n",
        "\n",
        "formatted_output = format_generated_output(before, after, output[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6zIQEruE5_FC"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "ORIGINAL PROMPT:\n",
            "import if __name__ == \"__main__\":\n",
            "    sys.exit(0)\n",
            "\n",
            "FORMATTED PROMPT:\n",
            "'<|fim_prefix|>import <|fim_suffix|>if __name__ == \"__main__\":\\n    sys.exit(0)<|fim_middle|>'\n",
            "\n",
            "GENERATED OUTPUT:\n",
            "'sys\\n<|file_separator|>'\n",
            "\n",
            "FILL-IN COMPLETION:\n",
            "import sys\n",
            "if __name__ == \"__main__\":\n",
            "    sys.exit(0)\n"
          ]
        }
      ],
      "source": [
        "before = \"import \"\n",
        "after = \"\"\"if __name__ == \"__main__\":\\n    sys.exit(0)\"\"\"\n",
        "prompt = format_completion_prompt(before, after)\n",
        "\n",
        "output = sampler(\n",
        "    [prompt],\n",
        "    total_generation_steps=100,\n",
        "    ).text\n",
        "\n",
        "formatted_output = format_generated_output(before, after, output[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SvaV4GU76M3t"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "ORIGINAL PROMPT:\n",
            "import numpy as np\n",
            "def reflect(matrix):\n",
            "  # horizontally reflect a matrix\n",
            "\n",
            "\n",
            "FORMATTED PROMPT:\n",
            "'<|fim_prefix|>import numpy as np\\ndef reflect(matrix):\\n  # horizontally reflect a matrix\\n<|fim_suffix|><|fim_middle|>'\n",
            "\n",
            "GENERATED OUTPUT:\n",
            "'  return np.flip(matrix, axis=1)\\n<|file_separator|>'\n",
            "\n",
            "FILL-IN COMPLETION:\n",
            "import numpy as np\n",
            "def reflect(matrix):\n",
            "  # horizontally reflect a matrix\n",
            "  return np.flip(matrix, axis=1)\n",
            "\n"
          ]
        }
      ],
      "source": [
        "before = \"\"\"import numpy as np\n",
        "def reflect(matrix):\n",
        "  # horizontally reflect a matrix\n",
        "\"\"\"\n",
        "after = \"\"\n",
        "prompt = format_completion_prompt(before, after)\n",
        "\n",
        "output = sampler(\n",
        "    [prompt],\n",
        "    total_generation_steps=100,\n",
        "    ).text\n",
        "\n",
        "formatted_output = format_generated_output(before, after, output[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jao0Qk-ZIqyD"
      },
      "source": [
        "## Learn more\n",
        "\n",
        "- You can learn more about the Google DeepMind [`gemma`  library on GitHub](https://github.com/google-deepmind/gemma), which contains docstrings of modules you used in this tutorial, such as [`gemma.params`](https://github.com/google-deepmind/gemma/blob/main/gemma/params.py),\n",
        "[`gemma.transformer`](https://github.com/google-deepmind/gemma/blob/main/gemma/transformer.py), and\n",
        "[`gemma.sampler`](https://github.com/google-deepmind/gemma/blob/main/gemma/sampler.py).\n",
        "- The following libraries have their own documentation sites: [core JAX](https://jax.readthedocs.io), [Flax](https://flax.readthedocs.io), and [Orbax](https://orbax.readthedocs.io/).\n",
        "- For `sentencepiece` tokenizer/detokenizer documentation, check out [Google's `sentencepiece` GitHub repo](https://github.com/google/sentencepiece).\n",
        "- For `kagglehub` documentation, check out `README.md` on [Kaggle's `kagglehub` GitHub repo](https://github.com/Kaggle/kagglehub).\n",
        "- Learn how to [use Gemma models with Google Cloud Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/open-models/use-gemma).\n",
        "- If you are using Google Cloud TPUs (v3-8 and newer), make sure to also update to the latest `jax[tpu]` package (`!pip install -U jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html`), restart the runtime, and check that `jax` and `jaxlib` versions match (`!pip list | grep jax`). This can prevent the `RuntimeError` that can arise because of the `jaxlib` and `jax` version mismatch. For more JAX installation instructions, refer to the [JAX docs](https://jax.readthedocs.io/en/latest/tutorials/installation.html#install-google-tpu)."
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "name": "codegemma_flax_inference.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
